#include <iostream>
#include <opencv2/core.hpp>
#include <opencv2/highgui.hpp>
#include <opencv2/imgproc.hpp>

#include "hi_ext_util.h"
#include "mpp_help.h"
#include "ai_plug.h"

#include "sample_comm_nnie.h"
#include "nnie_sample_plug.h"

using namespace std;
using namespace cv;

#define PLUG_UUID          "\"hi.tennis_detect\""
#define PLUG_DESC          "\"网球检测(传统算子)\""  // UTF8 encode

#define FRM_WIDTH          640
#define FRM_HEIGHT         480
#define TENNIS_OBJ_MAX     256
#define DRAW_RETC_THICK    2

static OsdSet* g_osdsTennis = NULL;
static int g_osd0Tennis = -1;

static const char TENNIS_DETECT[] = "{"
    "\"uuid\": " PLUG_UUID ","
    "\"desc\": " PLUG_DESC ","
    "\"frmWidth\": " HI_TO_STR(FRM_WIDTH) ","
    "\"frmHeight\": " HI_TO_STR(FRM_HEIGHT) ","
    "\"butt\": 0"
"}";

static const char* TennisDetectProf(void)
{
    return TENNIS_DETECT;
}

static int TennisDetectLoad(uintptr_t* model, OsdSet* osds)
{
    HI_S32 ret = 1;

    g_osdsTennis = osds;
    HI_ASSERT(g_osdsTennis);
    g_osd0Tennis = OsdsCreateRgn(g_osdsTennis);
    HI_ASSERT(g_osd0Tennis >= 0);
    *model = 1;
    LOGI("TennisDetectLoad success\n");
	
	return ret;
}

static int TennisDetectUnload(uintptr_t model)
{
    (void)model;
    OsdsClear(g_osdsTennis);
	
    return HI_SUCCESS;
}

typedef struct tagIPC_IMAGE{
    HI_U64 u64PhyAddr;
    HI_U64 u64VirAddr;
    HI_U32 u32Width;
    HI_U32 u32Height;
} IPC_IMAGE;

HI_S32 yuvFrame2rgb(VIDEO_FRAME_INFO_S *srcFrame, IPC_IMAGE *dstImage)
{
    IVE_HANDLE hIveHandle;
    IVE_SRC_IMAGE_S pstSrc;
    IVE_DST_IMAGE_S pstDst;
    IVE_CSC_CTRL_S stCscCtrl;
    HI_S32 s32Ret = 0;
    stCscCtrl.enMode = IVE_CSC_MODE_PIC_BT709_YUV2RGB; //IVE_CSC_MODE_VIDEO_BT601_YUV2RGB;
    //pstSrc初始化
    pstSrc.enType = IVE_IMAGE_TYPE_YUV420SP;
    pstSrc.au64VirAddr[0] = srcFrame->stVFrame.u64VirAddr[0];
    pstSrc.au64VirAddr[1] = srcFrame->stVFrame.u64VirAddr[1];
    pstSrc.au64VirAddr[2] = srcFrame->stVFrame.u64VirAddr[2];
 
    pstSrc.au64PhyAddr[0] = srcFrame->stVFrame.u64PhyAddr[0];
    pstSrc.au64PhyAddr[1] = srcFrame->stVFrame.u64PhyAddr[1];
    pstSrc.au64PhyAddr[2] = srcFrame->stVFrame.u64PhyAddr[2];
 
    pstSrc.au32Stride[0] = srcFrame->stVFrame.u32Stride[0];
    pstSrc.au32Stride[1] = srcFrame->stVFrame.u32Stride[1];
    pstSrc.au32Stride[2] = srcFrame->stVFrame.u32Stride[2];
 
    pstSrc.u32Width = srcFrame->stVFrame.u32Width;
    pstSrc.u32Height = srcFrame->stVFrame.u32Height;
    //ptrDst初始化
    pstDst.enType = IVE_IMAGE_TYPE_U8C3_PACKAGE;
    pstDst.u32Width = pstSrc.u32Width;
    pstDst.u32Height = pstSrc.u32Height;
    pstDst.au32Stride[0] = pstSrc.au32Stride[0];
    pstDst.au32Stride[1] = 0;
    pstDst.au32Stride[2] = 0;
    //pstDst虚拟地址分配
    s32Ret = HI_MPI_SYS_MmzAlloc_Cached(&pstDst.au64PhyAddr[0], (void **)&pstDst.au64VirAddr[0],
        "User", HI_NULL, pstDst.u32Height*pstDst.au32Stride[0] * 3);
    if (HI_SUCCESS != s32Ret) {       
        HI_MPI_SYS_MmzFree(pstDst.au64PhyAddr[0], (void *)pstDst.au64VirAddr[0]);
        LOGE("HI_MPI_SYS_MmzFree err\n");
        return s32Ret;
    }
    //FlushCache
    s32Ret = HI_MPI_SYS_MmzFlushCache(pstDst.au64PhyAddr[0], (void *)pstDst.au64VirAddr[0],
        pstDst.u32Height*pstDst.au32Stride[0]*3);
    if (HI_SUCCESS != s32Ret) {       
        HI_MPI_SYS_MmzFree(pstDst.au64PhyAddr[0], (void *)pstDst.au64VirAddr[0]);
        return s32Ret;
    }
    memset((void *)pstDst.au64VirAddr[0], 0, pstDst.u32Height*pstDst.au32Stride[0]*3);
    HI_BOOL bInstant = HI_TRUE;
    //色彩转换
    s32Ret = HI_MPI_IVE_CSC(&hIveHandle, &pstSrc, &pstDst, &stCscCtrl, bInstant);
    if(HI_SUCCESS != s32Ret) {       
        HI_MPI_SYS_MmzFree(pstDst.au64PhyAddr[0], (void *)pstDst.au64VirAddr[0]);
        return s32Ret;
    }
    //加急
    if (HI_TRUE == bInstant) {
        HI_BOOL bFinish = HI_TRUE;
        HI_BOOL bBlock = HI_TRUE;
        s32Ret = HI_MPI_IVE_Query(hIveHandle, &bFinish, bBlock);
        while (HI_ERR_IVE_QUERY_TIMEOUT == s32Ret) {
            usleep(100);
            s32Ret = HI_MPI_IVE_Query(hIveHandle,&bFinish,bBlock);
        }
    }
    dstImage->u64PhyAddr = pstDst.au64PhyAddr[0];
    dstImage->u64VirAddr = pstDst.au64VirAddr[0];
    dstImage->u32Width = pstDst.u32Width;
    dstImage->u32Height = pstDst.u32Height;
    
    return HI_SUCCESS;
}

HI_S32 frame2Mat(VIDEO_FRAME_INFO_S *srcFrame, Mat &dstMat)
{
    HI_U32 w = srcFrame->stVFrame.u32Width;
    HI_U32 h = srcFrame->stVFrame.u32Height;
    int bufLen = w * h * 3;
    HI_U8 *srcRGB = NULL;
    IPC_IMAGE dstImage;
    if (yuvFrame2rgb(srcFrame, &dstImage) != HI_SUCCESS) {
        LOGE("yuvFrame2rgb err\n");
        return HI_FAILURE;
    }
    srcRGB = (HI_U8 *)dstImage.u64VirAddr;
    dstMat.create(h, w, CV_8UC3);
    memcpy(dstMat.data, srcRGB, bufLen * sizeof(HI_U8));
    HI_MPI_SYS_MmzFree(dstImage.u64PhyAddr, (void *)&(dstImage.u64VirAddr));
    return HI_SUCCESS;
}

/**
    将计算结果打包为resJson.
*/
HI_CHAR* TennisDetectToJson(const RectBox items[], HI_S32 itemNum, int* resBytes)
{
    HI_S32 jsonSize = TINY_BUF_SIZE + itemNum * TINY_BUF_SIZE; // 每个item的打包size为TINY_BUF_SIZE
    HI_CHAR *jsonBuf = (HI_CHAR*)malloc(jsonSize);
    HI_ASSERT(jsonBuf);
    HI_S32 offset = 0;

    offset += snprintf_s(jsonBuf + offset, jsonSize - offset, jsonSize - offset - 1, "[");
    for (HI_S32 i = 0; i < itemNum; i++) {
        const RectBox *item = &items[i];
        offset += snprintf_s(jsonBuf + offset, jsonSize - offset, jsonSize - offset - 1,
            "%s { \"object xmin\": %d, \"ymin\": %d, \"xmax\": %d, \"ymax\": %d }",
            (i == 0 ? "\n  " : ", "), (uint)item->xmin, (uint)item->ymin, (uint)item->xmax, (uint)item->ymax);
        HI_ASSERT(offset < jsonSize);
    }
    offset += snprintf_s(jsonBuf + offset, jsonSize - offset, jsonSize - offset - 1, "]");
    HI_ASSERT(offset < jsonSize);
    
    if (resBytes) {
        *resBytes = offset;
    }
    return jsonBuf;
}

static int TennisDetectCal(uintptr_t model,
    VIDEO_FRAME_INFO_S *srcFrm, VIDEO_FRAME_INFO_S *dstFrm, HI_CHAR** resJson)
{
    (void)model;
    int ret = 0;
    RectBox boxs[TENNIS_OBJ_MAX] = {0};
    int j = 0;

    Mat image;
    frame2Mat(srcFrm, image);
    if (image.size == 0) {
        LOGD("image is null\n");
        return HI_FAILURE;
    }

    Mat src = image;
    Mat src1 = src.clone();
    Mat dst, edge, gray, hsv;

    dst.create(src1.size(), src1.type()); // 创建与src同类型和大小的矩阵(dst)
    // imwrite("image.jpg", src1);

    // cvtColor算子用于将图像从一个颜色空间转换到另一个颜色空间的转换
    cvtColor(src1, hsv, COLOR_BGR2HSV); // 将原图转换为HSV图像

    // 对hsv图像进行二值化处理，这里是将绿色背景二值化，该参数根据需求进行调整
    inRange(hsv, Scalar(31, 82, 68), Scalar(65, 248, 255), gray);

    // 利用canny算子进行边缘检测
    Canny(gray, gray, 3, 9, 3);
    vector<vector<Point>> contours;
    findContours(gray, contours, RETR_EXTERNAL, CHAIN_APPROX_SIMPLE, Point());
    LOGI("contours.size():%d\n", contours.size());

    for (int i = 0; i < (int)contours.size(); i++) {
        if (contours.size() > 30) {
            continue;
        }

        Rect ret1 = boundingRect(Mat(contours[i]));
        ret1.x -= 5;
        ret1.y -= 5;
        ret1.width += 10;
        ret1.height += 10;

        if ((ret1.width > 20) && (ret1.height > 20)) {
            LOGD("ret x:%d, y:%d, width:%d, height:%d\n", ret1.x, ret1.y, ret1.width, ret1.height);
            boxs[j].xmin = ret1.x * 2;
            boxs[j].ymin = (int)(ret1.y * 1.5);
            boxs[j].xmax = boxs[j].xmin + ret1.width * 2;
            boxs[j].ymax = boxs[j].ymin + (int)ret1.height * 1.5;
            j++;
        }
        LOGD("dstfrm width:%d, Height:%d\n", dstFrm->stVFrame.u32Width, dstFrm->stVFrame.u32Height);
	}

    // 打包计算结果为resJson
    *resJson = TennisDetectToJson(boxs, j, NULL);

    if (j > 0 && j <= 25) {
        LOGI("box num:%d\n", j);
        MppFrmDrawRects(dstFrm, boxs, j, RGB888_RED, DRAW_RETC_THICK);
    }

    return ret;
}

static const AiPlug G_TENNIS_DETECT_ITF = {
    .Prof = TennisDetectProf,
    .Load = TennisDetectLoad,
    .Unload = TennisDetectUnload,
    .Cal = TennisDetectCal,
};

const AiPlug* AiPlugItf(uint32_t* magic)
{
    if (magic) {
        *magic = AI_PLUG_MAGIC;
    }

    return (AiPlug*)&G_TENNIS_DETECT_ITF;
}
